#include <openssl/ec.h>
#include <openssl/evp.h>
#include <openssl/objects.h>
#include <openssl/err.h>
#include "sm2.h"
#include <stdio.h>
#include <stdlib.h>

EC_KEY* generate_sm2_keypair() {
    EC_KEY *eckey = EC_KEY_new_by_curve_name(NID_sm2);
    if (!eckey) {
        ERR_print_errors_fp(stderr);
        return NULL;
    }
    if (!EC_KEY_generate_key(eckey)) {
        EC_KEY_free(eckey);
        ERR_print_errors_fp(stderr);
        return NULL;
    }
    return eckey;
}

int sm2_encrypt(const unsigned char *plaintext, size_t plaintext_len, EC_KEY *pub_key, unsigned char **ciphertext, size_t *ciphertext_len) {
    EVP_PKEY *pkey = EVP_PKEY_new();
    if (!pkey) {
        ERR_print_errors_fp(stderr);
        return 0;
    }
    EVP_PKEY_set1_EC_KEY(pkey, pub_key);

    EVP_PKEY_CTX *ctx = EVP_PKEY_CTX_new(pkey, NULL);
    if (!ctx) {
        EVP_PKEY_free(pkey);
        ERR_print_errors_fp(stderr);
        return 0;
    }

    if (EVP_PKEY_encrypt_init(ctx) <= 0) {
        EVP_PKEY_CTX_free(ctx);
        EVP_PKEY_free(pkey);
        ERR_print_errors_fp(stderr);
        return 0;
    }

    if (EVP_PKEY_encrypt(ctx, NULL, ciphertext_len, plaintext, plaintext_len) <= 0) {
        EVP_PKEY_CTX_free(ctx);
        EVP_PKEY_free(pkey);
        ERR_print_errors_fp(stderr);
        return 0;
    }

    *ciphertext = OPENSSL_malloc(*ciphertext_len);
    if (!*ciphertext) {
        EVP_PKEY_CTX_free(ctx);
        EVP_PKEY_free(pkey);
        ERR_print_errors_fp(stderr);
        return 0;
    }

    if (EVP_PKEY_encrypt(ctx, *ciphertext, ciphertext_len, plaintext, plaintext_len) <= 0) {
        OPENSSL_free(*ciphertext);
        *ciphertext = NULL;
        EVP_PKEY_CTX_free(ctx);
        EVP_PKEY_free(pkey);
        ERR_print_errors_fp(stderr);
        return 0;
    }

    EVP_PKEY_CTX_free(ctx);
    EVP_PKEY_free(pkey);
    return 1;
}

int sm2_decrypt(const unsigned char *ciphertext, size_t ciphertext_len, EC_KEY *priv_key, unsigned char **plaintext, size_t *plaintext_len) {
    EVP_PKEY *pkey = EVP_PKEY_new();
    if (!pkey) {
        ERR_print_errors_fp(stderr);
        return 0;
    }
    EVP_PKEY_set1_EC_KEY(pkey, priv_key);

    EVP_PKEY_CTX *ctx = EVP_PKEY_CTX_new(pkey, NULL);
    if (!ctx) {
        EVP_PKEY_free(pkey);
        ERR_print_errors_fp(stderr);
        return 0;
    }

    if (EVP_PKEY_decrypt_init(ctx) <= 0) {
        EVP_PKEY_CTX_free(ctx);
        EVP_PKEY_free(pkey);
        ERR_print_errors_fp(stderr);
        return 0;
    }

    if (EVP_PKEY_decrypt(ctx, NULL, plaintext_len, ciphertext, ciphertext_len) <= 0) {
        EVP_PKEY_CTX_free(ctx);
        EVP_PKEY_free(pkey);
        ERR_print_errors_fp(stderr);
        return 0;
    }

    *plaintext = OPENSSL_malloc(*plaintext_len);
    if (!*plaintext) {
        EVP_PKEY_CTX_free(ctx);
        EVP_PKEY_free(pkey);
        ERR_print_errors_fp(stderr);
        return 0;
    }

    if (EVP_PKEY_decrypt(ctx, *plaintext, plaintext_len, ciphertext, ciphertext_len) <= 0) {
        OPENSSL_free(*plaintext);
        *plaintext = NULL;
        EVP_PKEY_CTX_free(ctx);
        EVP_PKEY_free(pkey);
        ERR_print_errors_fp(stderr);
        return 0;
    }

    EVP_PKEY_CTX_free(ctx);
    EVP_PKEY_free(pkey);
    return 1;
}
